import numpy as np
import pandas as pd


class Homography_8:

    get_x = lambda x, y, H: (x * H[0, 0] + y * H[0, 1] + H[0, 2])/(x * H[2, 0] + y * H[2, 1] + 1.0)
    get_y = lambda x, y, H: (x * H[1, 0] + y * H[1, 1] + H[1, 2])/(x * H[2, 0] + y * H[2, 1] + 1.0)

    class error(Exception):
        def __init__(self, message):
            """
            自定义异常
            :param message: 抛出异常时的提示
            """
            self.message = message
            super().__init__(self.message)

    def __init__(self, pts_reg: np.ndarray, pts_ref: np.ndarray) -> None:
        if pts_reg.shape[0] != pts_ref.shape[0]:
            raise Homography_8.error('配准点和参考点的数量不等')
        if pts_reg.shape[0] == 0:
            raise Homography_8.error('没有设置特征点')
        if pts_reg.shape[1] != 2:
            raise Homography_8.error('配准点必须是二维的')
        if pts_ref.shape[1] != 2:
            raise Homography_8.error('参考点必须是二维的')
        self.num = pts_reg.shape[0]
        self.pts_reg = pts_reg
        self.pts_ref = pts_ref

    def get_homography(self):
        """
        解析单应性矩阵
        :return:
        """
        import cv2 as cv
        if self.num < 4:
            raise Homography_8.error('配准点的数量必须大于等于4')
        H, _ = cv.findHomography(self.pts_reg, self.pts_ref)
        self.H = H

    def get_error(self) -> pd.DataFrame:
        """
        计算配准误差
        :return: 返回误差
        """
        result = []
        for i in range(self.num):
            try:
                trans_x = Homography_8.get_x(self.pts_reg[i, 0], self.pts_reg[i, 1], self.H)
                trans_y = Homography_8.get_y(self.pts_reg[i, 0], self.pts_reg[i, 1], self.H)
                error_x = trans_x - self.pts_ref[i, 0]
                error_y = trans_y - self.pts_ref[i, 1]
            except TypeError:
                trans_x = trans_y = error_x = error_y = None
            result.append([self.pts_reg[i, 0], self.pts_reg[i, 1], trans_x, trans_y,
                           self.pts_ref[i, 0], self.pts_ref[i, 1], error_x, error_y])
        result: pd.DataFrame = pd.DataFrame(result, 
                                            columns=['reg_x', 'reg_y', 'trans_x', 'trans_y', 'ref_x', 'ref_y', 'error_x', 'error_y'])
        return result


class Homography_6(Homography_8):
    def __init__(self, pts_reg: np.ndarray, pts_ref: np.ndarray) -> None:
        super().__init__(pts_reg, pts_ref)

    def get_homography(self):
        """
        计算自由度为6的单应性矩阵
        :return:
        """
        if self.num < 3:
            raise Homography_8.error('配准点的数量必须大于等于3')
        num = self.num
        A = np.zeros([2 * num, 6], dtype=np.float64)
        b = np.zeros([2 * num], dtype=np.float64)
        for i in range(num):
            A[2 * i, 0] = self.pts_reg[i, 0]
            A[2 * i, 1] = self.pts_reg[i, 1]
            A[2 * i, 2] = 1.0
            b[2 * i] = self.pts_ref[i, 0]
            A[2 * i + 1, 3] = self.pts_reg[i, 0]
            A[2 * i + 1, 4] = self.pts_reg[i, 1]
            A[2 * i + 1, 5] = 1.0
            b[2 * i + 1] = self.pts_ref[i, 1]
        result = np.linalg.lstsq(A, b)
        h = result[0]
        H = np.zeros([3, 3], dtype=np.float64)
        H[0, :] = h[: 3]
        H[1, :] = h[3:]
        H[2, 2] = 1.0
        self.H = H


class Homography_4(Homography_8):
    def __init__(self, pts_reg: np.ndarray, pts_ref: np.ndarray) -> None:
        super().__init__(pts_reg, pts_ref)
    
    def get_homography(self):
        """
        解析配准量，包括缩放量、旋转量（角度制）、平移量，并计算单应性矩阵
        :return:
        """
        from numpy import cos, sin
        if self.num < 2:
            raise Homography_8.error('配准点的数量必须大于等于2')
        # 计算两组坐标的中心点
        center_reg = np.mean(self.pts_reg, axis=0)
        center_ref = np.mean(self.pts_ref, axis=0)
        # 去中心化
        decentered_reg = self.pts_reg - center_reg
        decentered_ref = self.pts_ref - center_ref
        # 将去中心后的坐标转换为复数形式
        complex_reg = decentered_reg[:, 0] + 1j * decentered_reg[:, 1]
        complex_ref = decentered_ref[:, 0] + 1j * decentered_ref[:, 1]
        # 计算两组复数序列的离散傅里叶变换（DFT）
        fft_reg = np.fft.fft(complex_reg)
        fft_ref = np.fft.fft(complex_ref)
        fft_reg_conj = np.conjugate(fft_reg)  # 计算两组DFT结果的共轭
        cross_power_spectrum = np.multiply(fft_reg_conj, fft_ref)  # 计算互功率谱（点乘共轭）
        # cross_power_spectrum = np.divide(cross_power_spectrum, np.abs(fft_reg))
        # cross_power_spectrum = np.divide(cross_power_spectrum, np.abs(fft_reg))
        index_max = np.argmax(np.abs(cross_power_spectrum))  # 找到相位相关峰值的位置，这将是位移的估计
        angle = np.angle(cross_power_spectrum[index_max], deg=False)  # 计算相位（弧度制）
        scale = np.abs(cross_power_spectrum[index_max]) / np.abs(fft_reg[index_max]) ** 2
        shift_x = -scale * cos(angle) * center_reg[0] + scale * sin(angle) * center_reg[1] + center_ref[0]
        shift_y = -scale * sin(angle) * center_reg[0] - scale * cos(angle) * center_reg[1] + center_ref[1]
        H = np.array([scale * cos(angle), -scale * sin(angle), shift_x,
                      scale * sin(angle), scale * cos(angle), shift_y,
                      0, 0, 1.0])
        angle = np.rad2deg(angle)
        angle %= 360
        self.angle = angle
        self.scale = scale
        self.shift_x = shift_x
        self.shift_y = shift_y
        self.H = H.reshape([3, 3])


def set_homography(arg: np.ndarray):
    """
    根据参数设置单应性矩阵。（自由度分别为4/6/8）
    :param arg: 生成单应性矩阵的参数
    :return:
    """
    from numpy import cos, sin
    if arg.shape[0] == 4:
        scale = arg[0]
        angle = arg[1]
        angle = np.deg2rad(angle)
        shift_x = arg[2]
        shift_y = arg[3]
        H = np.array([scale * cos(angle), -scale * sin(angle), shift_x,
                      scale * sin(angle), scale * cos(angle), shift_y,
                      0, 0, 1.0])
        H = H.reshape([3, 3])
    if arg.shape[0] == 6:
        H = np.zeros([3, 3], dtype=np.float64)
        H[0, :] = arg[: 3]
        H[1, :] = arg[3:]
        H[2, 2] = 1.0
    if arg.shape[0] == 8:
        H = np.zeros([3, 3], dtype=np.float64)
        H[0, :] = arg[: 3]
        H[1, :] = arg[3: 6]
        H[2, : 2] = arg[6:]
        H[2, 2] = 1.0
    return H


def generate_pt(num: int, H: np.ndarray, is_noise=True):
    """
    生成配准点和参考点，默认加入噪声
    :param num: 随机点的数量
    :param H: 单应性矩阵
    :param is_noise: 是否加入噪声
    :return:
    """
    pt_reg = np.random.uniform(0, 1000, (num, 2))
    pt_ref = np.zeros_like(pt_reg)
    for i in range(num):
        pt_ref[i, 0] = Homography_8.get_x(pt_reg[i, 0], pt_reg[i, 1], H)
        pt_ref[i, 1] = Homography_8.get_y(pt_reg[i, 0], pt_reg[i, 1], H)
    if is_noise:
        noise = np.random.normal(0, 2, pt_ref.shape)
        pt_ref += noise
    return pt_reg, pt_ref


if __name__ == '__main__':
    arg = np.array([2, 30, 12, 16])  # 幅度、角度、平移量
    print(arg)
    H_set = set_homography(arg)
    print(H_set)
    pt_reg, pt_ref = generate_pt(8, H_set, is_noise=True)
    registration = Homography_4(pt_reg, pt_ref)
    registration.get_homography()
    result = registration.get_error()
    print(result)
    print(registration.scale, registration.angle, registration.shift_x, registration.shift_y)
    print(registration.H)

    # arg = np.array([2.1, 2.6, 8.1, 2.3, 2.5, 7.4])
    # print(arg)
    # H_set = set_homography(arg)
    # print(H_set)
    # pt_reg, pt_ref = generate_pt(8, H_set, is_noise=True)
    # registration = Homography_6(pt_reg, pt_ref)
    # registration.get_homography()
    # result = registration.get_error()
    # print(result)
    # print(registration.H)

    # arg = np.array([2.1, 2.6, 8.1, 2.3, 2.5, 7.4, 1e-3, -2e-3])
    # print(arg)
    # H_set = set_homography(arg)
    # print(H_set)
    # pt_reg, pt_ref = generate_pt(8, H_set, is_noise=False)
    # registration = Homography(pt_reg, pt_ref)
    # registration.get_homography()
    # result = registration.get_error()
    # print(result)
    # print(registration.H)
    
    pass
